import os
import numpy as np
import math
from scipy.spatial.transform import Rotation as R
import json

'''
origin author: 水木皆Ming
重力补偿计算
'''


class GravityCompensation:
    '''
    计算负载重力、质心、实际所受外力
    数学推导：https://blog.csdn.net/shuimujieming/article/details/137015352?spm=1001.2014.3001.5501
             https://kns.cnki.net/kcms2/article/abstract?v=Uq4Diyda8XBnoha3opN76DCc52lWeFtkUg32VJKUCarSkF9_iOzDTbp_97fL8v0fkfvXlMGxZeeEdgM4HBFvglPqZbc3S6ZQ4ChDLJA2Nl8A0cQdtWZu6eCxl2DiCvUPRq9V2gbrGgNGL8ea5tY0DhckAzNQsTGguggWyv_LzBfntJl1hEqhPfa7IanMNcKbM1HUCk7shoMhECxuINr4eD9UdpNaqHJfTbkFTcN9Jzc=&uniplatform=NZKPT
    '''

    def __init__(self):
        # 存储多个点位的数据
        self.M = np.empty((0, 0))  # 力矩
        self.F = np.empty((0, 0))  # 力
        self.f = np.empty((0, 0))  # 力，和上面那个存储格式不一样
        self.R = np.empty((0, 0))  # 末端到基坐标的旋转矩阵

        self.x = 0
        self.y = 0
        self.z = 0
        self.k1 = 0
        self.k2 = 0
        self.k3 = 0

        self.U = 0  # 机器人安装角度
        self.V = 0  # 机器人安装角度
        self.g = 0  # 负载重力

        # 力分量的零点值
        self.F_x0 = 0
        self.F_y0 = 0
        self.F_z0 = 0
        # 力矩分量的零点值
        self.M_x0 = 0
        self.M_y0 = 0
        self.M_z0 = 0

        self.load_config()

    def Update_M(self, torque_data):
        M_x = torque_data[0]
        M_y = torque_data[1]
        M_z = torque_data[2]

        if (self.M.any()):
            M_1 = np.matrix([M_x, M_y, M_z]).transpose()
            self.M = np.vstack((self.M, M_1))
        else:
            self.M = np.matrix([M_x, M_y, M_z]).transpose()

    def Update_F(self, force_data):
        F_x = force_data[0]
        F_y = force_data[1]
        F_z = force_data[2]

        if (self.F.any()):
            F_1 = np.matrix([[0, F_z, -F_y, 1, 0, 0],
                             [-F_z, 0, F_x, 0, 1, 0],
                             [F_y, -F_x, 0, 0, 0, 1]])
            self.F = np.vstack((self.F, F_1))
        else:
            self.F = np.matrix([[0, F_z, -F_y, 1, 0, 0],
                                [-F_z, 0, F_x, 0, 1, 0],
                                [F_y, -F_x, 0, 0, 0, 1]])

    def Solve_A(self):
        A = np.dot(np.dot(np.linalg.inv(np.dot(self.F.transpose(), self.F)), self.F.transpose()), self.M)

        self.x = A[0, 0]
        self.y = A[1, 0]
        self.z = A[2, 0]
        self.k1 = A[3, 0]
        self.k2 = A[4, 0]
        self.k3 = A[5, 0]
        # print("A= \n" , A)
        print("x= ", self.x)
        print("y= ", self.y)
        print("z= ", self.z)
        print("k1= ", self.k1)
        print("k2= ", self.k2)
        print("k3= ", self.k3)

    def Update_f(self, force_data):
        F_x = force_data[0]
        F_y = force_data[1]
        F_z = force_data[2]

        if (self.f.any()):
            f_1 = np.matrix([F_x, F_y, F_z]).transpose()
            self.f = np.vstack((self.f, f_1))
        else:
            self.f = np.matrix([F_x, F_y, F_z]).transpose()

    def Update_R(self, rotation_vector):
        # 机械臂末端到基坐标的旋转矩阵
        # R_array = self.eulerAngles2rotationMat(euler_data)
        rotation = R.from_rotvec(rotation_vector)
        R_array = rotation.as_matrix()

        alpha = (0) * 180 / np.pi

        # 力传感器到末端的旋转矩阵
        R_alpha = np.array([[math.cos(alpha), -math.sin(alpha), 0],
                            [math.sin(alpha), math.cos(alpha), 0],
                            [0, 0, 1]
                            ])

        R_array = np.dot(R_alpha, R_array.transpose())

        if (self.R.any()):
            R_1 = np.hstack((R_array, np.eye(3)))
            self.R = np.vstack((self.R, R_1))
        else:
            self.R = np.hstack((R_array, np.eye(3)))

    def Solve_B(self):
        # 必须在Solve_A之后调用
        B = np.dot(np.dot(np.linalg.inv(np.dot(self.R.transpose(), self.R)), self.R.transpose()), self.f)

        self.g = math.sqrt(B[0] * B[0] + B[1] * B[1] + B[2] * B[2])
        self.U = math.asin(-B[1] / self.g)
        self.V = math.atan(-B[0] / B[2])

        self.F_x0 = B[3, 0]
        self.F_y0 = B[4, 0]
        self.F_z0 = B[5, 0]

        self.M_x0 = self.k1 - self.F_y0 * self.z + self.F_z0 * self.y
        self.M_y0 = self.k2 - self.F_z0 * self.x + self.F_x0 * self.z
        self.M_z0 = self.k3 - self.F_x0 * self.y + self.F_y0 * self.x

        # print("B= \n" , B)
        print("g= ", self.g / 9.81)
        print("U= ", self.U * 180 / math.pi)
        print("V= ", self.V * 180 / math.pi)
        print("F_x0= ", self.F_x0)
        print("F_y0= ", self.F_y0)
        print("F_z0= ", self.F_z0)
        print("M_x0= ", self.M_x0)
        print("M_y0= ", self.M_y0)
        print("M_z0= ", self.M_z0)

    def Solve_Force(self, force_data, rotation_vector):
        Force_input = np.matrix([force_data[0], force_data[1], force_data[2]]).transpose()

        my_f = np.matrix([np.cos(self.U) * np.sin(self.V) * self.g, -np.sin(self.U) * self.g,
                          -np.cos(self.U) * np.cos(self.V) * self.g, self.F_x0, self.F_y0, self.F_z0]).transpose()

        rotation = R.from_rotvec(rotation_vector)
        R_array = rotation.as_matrix()
        R_array = R_array.transpose()
        R_1 = np.hstack((R_array, np.eye(3)))

        Force_ex = Force_input - np.dot(R_1, my_f)
        # print('接触力：', Force_ex.T)
        return Force_ex

    def Solve_Torque(self, torque_data, rotation_vector):
        Torque_input = np.matrix([torque_data[0], torque_data[1], torque_data[2]]).transpose()

        Torque_zero = np.matrix([self.M_x0, self.M_y0, self.M_z0]).transpose()

        Gravity_param = np.matrix([[0, -self.z, self.y],
                                   [self.z, 0, -self.x],
                                   [-self.y, self.x, 0]])

        Gravity_input = np.matrix([np.cos(self.U) * np.sin(self.V) * self.g, -np.sin(self.U) * self.g,
                                   -np.cos(self.U) * np.cos(self.V) * self.g]).transpose()

        rotation = R.from_rotvec(rotation_vector)
        R_array = rotation.as_matrix()
        R_array = R_array.transpose()

        Torque_ex = Torque_input - Torque_zero - np.dot(np.dot(Gravity_param, R_array), Gravity_input)
        return Torque_ex
        # print('接触力矩：', Torque_ex.T)

    def GetForceTorque(self, force_torque_data, rotation_vector) -> list:
        '''
        计算传感器数据在传感器坐标系下的力和力矩
        Args:
            force_torque_data: 传感器数据
            rotation_vector: 传感器(机械臂末端)到基坐标的旋转向量
        '''
        force_data = force_torque_data[0:3]
        torque_data = force_torque_data[3:6]

        force = self.Solve_Force(force_data, rotation_vector)
        torque = self.Solve_Torque(torque_data, rotation_vector)

        force_list = np.asarray(force).flatten().tolist()
        torque_list = np.asarray(torque).flatten().tolist()

        return force_list + torque_list

    def GetForceTorque_On_Base(self, force_torque_data, rotation_vector) -> list:
        '''
        计算传感器数据在基坐标系下的力和力矩
        Args:
            force_torque_data: 传感器数据
            rotation_vector: 传感器(机械臂末端)到基坐标的旋转向量
        '''
        ft_on_sensor = self.GetForceTorque(force_torque_data, rotation_vector)
        force_on_sensor = np.array(ft_on_sensor[0:3])
        torque_on_sensor = np.array(ft_on_sensor[3:6])
        rotation = R.from_rotvec(rotation_vector)
        R_array = rotation.as_matrix()
        R_array_inv = R_array.transpose()  # 旋转矩阵的逆矩阵等于转置矩阵

        # force_on_base = np.dot(R_array_inv, force_on_sensor)
        # torque_on_base =np.dot(R_array_inv, torque_on_sensor)

        force_on_base = np.dot(force_on_sensor, R_array_inv)
        torque_on_base = np.dot(torque_on_sensor, R_array_inv)

        force_on_base_list = force_on_base.tolist()
        torque_on_base_list = torque_on_base.tolist()

        return force_on_base_list + torque_on_base_list

    def FtZero(self, force_torque_data, rotation_vector):
        force_torque = self.GetForceTorque(force_torque_data, rotation_vector)
        self.F_x0 += force_torque[0]
        self.F_y0 += force_torque[1]
        self.F_z0 += force_torque[2]
        self.M_x0 += force_torque[3]
        self.M_y0 += force_torque[4]
        self.M_z0 += force_torque[5]
        self.save_config()

    def save_config(self, filename="config.json"):
        # 保存类中的参数到JSON文件
        config = {
            'x': self.x,
            'y': self.y,
            'z': self.z,
            'k1': self.k1,
            'k2': self.k2,
            'k3': self.k3,
            'U': self.U,
            'V': self.V,
            'g': self.g,
            'F_x0': self.F_x0,
            'F_y0': self.F_y0,
            'F_z0': self.F_z0,
            'M_x0': self.M_x0,
            'M_y0': self.M_y0,
            'M_z0': self.M_z0
        }
        with open(filename, 'w') as f:
            json.dump(config, f, indent=4)
        print(f"重力补偿参数已保存到 {filename}。")

    def load_config(self, filename="config.json"):
        # 如果配置文件存在，读取文件并更新类中的参数
        if os.path.exists(filename):
            with open(filename, 'r') as f:
                config = json.load(f)
                self.x = config['x']
                self.y = config['y']
                self.z = config['z']
                self.k1 = config['k1']
                self.k2 = config['k2']
                self.k3 = config['k3']
                self.U = config['U']
                self.V = config['V']
                self.g = config['g']
                self.F_x0 = config['F_x0']
                self.F_y0 = config['F_y0']
                self.F_z0 = config['F_z0']
                self.M_x0 = config['M_x0']
                self.M_y0 = config['M_y0']
                self.M_z0 = config['M_z0']
        else:
            print(f"重力补偿配置文件 {filename} 不存在，使用默认参数。")

    def eulerAngles2rotationMat(self, theta):
        theta = [i * math.pi / 180.0 for i in theta]  # 角度转弧度

        R_x = np.array([[1, 0, 0],
                        [0, math.cos(theta[0]), -math.sin(theta[0])],
                        [0, math.sin(theta[0]), math.cos(theta[0])]
                        ])

        R_y = np.array([[math.cos(theta[1]), 0, math.sin(theta[1])],
                        [0, 1, 0],
                        [-math.sin(theta[1]), 0, math.cos(theta[1])]
                        ])

        R_z = np.array([[math.cos(theta[2]), -math.sin(theta[2]), 0],
                        [math.sin(theta[2]), math.cos(theta[2]), 0],
                        [0, 0, 1]
                        ])

        # 第一个角为绕X轴旋转，第二个角为绕Y轴旋转，第三个角为绕Z轴旋转
        R = np.dot(R_x, np.dot(R_y, R_z))
        return R


def run():
    import rtde_receive
    import rtde_control
    from Hardware.KunWei_ftsensor import ftsensor
    # from Hardware.FT300 import RobotiqFt300
    import time

    ip = '192.168.68.250'
    rtde_r = rtde_receive.RTDEReceiveInterface(ip)
    rtde_c = rtde_control.RTDEControlInterface(ip)
    flange_relative_pos = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    ft = ftsensor(False)
    # ft=RobotiqFt300()
    compensation = GravityCompensation()

    # UR5
    # collect_data_pos = [[np.pi,-np.pi/2,-np.pi/2,-np.pi/2,-3*np.pi/2,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-np.pi/2,-5*np.pi/4,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-np.pi/2,-3*np.pi/4,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi/4,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-np.pi/2,0,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-3*np.pi/4,-np.pi/2,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-np.pi,-np.pi/2,0],
    #                     [np.pi,-np.pi/2,-np.pi/2,-5*np.pi/4,-np.pi/2,0]]

    # UR5e
    # collect_data_pos = [[-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,np.pi/2,-np.pi],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, -3*np.pi/4],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi/2],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi/4],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,np.pi/2,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,np.pi/4,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,0,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi/4,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-3*np.pi/4,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-3*np.pi/4,-np.pi/2,0],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi,-np.pi/2,0],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -np.pi/4],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -np.pi / 2],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -3*np.pi / 4],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -np.pi],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi/4*3],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi/2],
    #                     [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi/4],
    #                     [-3*np.pi/2,-np.pi/2,-np.pi/2,-5*np.pi/4,-np.pi/2,0]]

    collect_data_pos = [
        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, 0],
        [-3 * np.pi / 2, -105.73 / 180 * np.pi, -87.5 / 180 * np.pi, -39.02 / 180 * np.pi, 100.48 / 180 * np.pi, 0],
        [-3 * np.pi / 2, -113.75 / 180 * np.pi, -87.63 / 180 * np.pi, 42.5 / 180 * np.pi, 100.44 / 180 * np.pi, 0],
        [-3 * np.pi / 2, -114.15 / 180 * np.pi, -87.91 / 180 * np.pi, -22.08 / 180 * np.pi, 50.21 / 180 * np.pi, 0]
    ]

    print('start')
    print('##############################################')

    # pos=rtde_r.getActualTCPPose()
    # if pos[0]>0:
    #     rtde_c.moveJ([0,-np.pi/2,0,-np.pi/2,-3*np.pi/2,0],0.2,0.05)
    #     pos=rtde_r.getActualTCPPose()
    #     pos[0]-=0.15
    #     rtde_c.moveL(pos,0.2,0.05)
    #     joint=rtde_r.getActualQ()
    #     joint[0]+=np.pi/4
    #     rtde_c.moveJ(joint,0.2,0.05)
    #     joint=rtde_r.getActualQ()
    #     joint[0]+=np.pi/4
    #     rtde_c.moveJ(joint,0.2,0.05)
    #     joint=rtde_r.getActualQ()
    #     joint[0]+=np.pi/4
    #     rtde_c.moveJ(joint,0.2,0.05)

    for i in range(len(collect_data_pos)):
        rtde_c.moveJ(collect_data_pos[i], 0.2, 0.05)
        time.sleep(1)
        force_torque = ft.GetForce()
        # force_torque = rtde_r.getFtRawWrench()
        force_data = force_torque[0:3]
        torque_data = force_torque[3:6]
        flange_pos = rtde_c.getForwardKinematics(collect_data_pos[i], flange_relative_pos)
        rotation_vector = flange_pos[3:6]
        compensation.Update_F(force_data)
        compensation.Update_M(torque_data)
        compensation.Update_f(force_data)
        compensation.Update_R(rotation_vector)
        print('第', i + 1, '次数据记录完成')

    print('数据记录完成，开始计算补偿参数')
    print('##############################################')
    compensation.Solve_A()
    compensation.Solve_B()
    print('补偿参数计算完成')
    compensation.save_config()
    ft.DisConnect()


if __name__ == '__main__':
    run()
